#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
微信公众号文章信息提取器
用于批量提取公众号文章并生成Excel报告
"""

import json
import time
import pandas as pd
from datetime import datetime
from urllib.parse import urlencode
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.options import Options
from selenium.common.exceptions import TimeoutException
import requests
from bs4 import BeautifulSoup

from config import WECHAT_CONFIG, DATE_RANGE, OUTPUT_CONFIG, BROWSER_CONFIG


class WeChatArticleExtractor:
    """微信公众号文章提取器类"""
    
    def __init__(self):
        self.driver = None
        self.articles = []
        self.session = requests.Session()
        
    def setup_browser(self):
        """设置浏览器"""
        chrome_options = Options()
        if BROWSER_CONFIG['headless']:
            chrome_options.add_argument('--headless')
        chrome_options.add_argument('--no-sandbox')
        chrome_options.add_argument('--disable-dev-shm-usage')
        chrome_options.add_argument(f'--window-size={BROWSER_CONFIG["window_size"]}')
        chrome_options.add_argument(f'--user-agent={BROWSER_CONFIG["user_agent"]}')
        
        self.driver = webdriver.Chrome(options=chrome_options)
        self.driver.implicitly_wait(10)
        
    def login(self, username, password):
        """登录微信公众号后台"""
        print("正在登录微信公众号后台...")
        self.driver.get(WECHAT_CONFIG['login_url'])
        
        try:
            # 等待登录页面加载
            WebDriverWait(self.driver, 20).until(
                EC.presence_of_element_located((By.NAME, "account"))
            )
            
            # 输入用户名和密码
            username_input = self.driver.find_element(By.NAME, "account")
            password_input = self.driver.find_element(By.NAME, "password")
            
            username_input.clear()
            username_input.send_keys(username)
            password_input.clear()
            password_input.send_keys(password)
            
            # 点击登录按钮
            login_button = self.driver.find_element(By.CLASS_NAME, "btn_login")
            login_button.click()
            
            # 等待登录成功
            WebDriverWait(self.driver, 30).until(
                EC.url_contains("cgi-bin/home")
            )
            print("登录成功！")
            return True
            
        except TimeoutException:
            print("登录超时，请检查用户名密码是否正确")
            return False
    
    def get_article_list(self, fakeid, token, begin=0, count=5):
        """获取文章列表"""
        url = WECHAT_CONFIG['article_list_url']
        
        params = {
            'token': token,
            'lang': 'zh_CN',
            'f': 'json',
            'ajax': '1',
            'random': '0.123456789',
            'action': 'list_ex',
            'begin': str(begin),
            'count': str(count),
            'query': '',
            'fakeid': fakeid,
            'type': '9'
        }
        
        headers = {
            'User-Agent': BROWSER_CONFIG['user_agent'],
            'Referer': 'https://mp.weixin.qq.com/cgi-bin/appmsg?t=media/appmsg_list&type=9&action=list&lang=zh_CN&token=' + token
        }
        
        try:
            response = self.session.get(url, params=params, headers=headers, timeout=WECHAT_CONFIG['timeout'])
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"获取文章列表失败: {e}")
            return None
    
    def parse_articles(self, data):
        """解析文章数据"""
        if not data or 'app_msg_list' not in data:
            return []
        
        articles = []
        for item in data['app_msg_list']:
            try:
                # 解析发布时间
                publish_time = datetime.fromtimestamp(item['update_time'])
                
                # 检查时间范围
                start_date = datetime.strptime(DATE_RANGE['start_date'], '%Y-%m-%d')
                end_date = datetime.strptime(DATE_RANGE['end_date'], '%Y-%m-%d')
                
                if start_date <= publish_time <= end_date:
                    article = {
                        'title': item['title'],
                        'publish_time': publish_time.strftime('%Y-%m-%d %H:%M:%S'),
                        'location': '微信公众号',  # 默认地点
                        'url': item['link']
                    }
                    articles.append(article)
                    
            except Exception as e:
                print(f"解析文章数据失败: {e}")
                continue
                
        return articles
    
    def extract_all_articles(self, fakeid, token):
        """提取所有符合条件的文章"""
        print("开始提取文章...")
        all_articles = []
        begin = 0
        count = 5
        
        while True:
            print(f"正在获取第 {begin//count + 1} 页文章...")
            data = self.get_article_list(fakeid, token, begin, count)
            
            if not data or 'app_msg_list' not in data:
                break
                
            articles = self.parse_articles(data)
            if not articles:
                break
                
            all_articles.extend(articles)
            begin += count
            
            # 避免请求过快
            time.sleep(2)
        
        self.articles = all_articles
        print(f"共提取到 {len(all_articles)} 篇符合条件的文章")
    
    def get_token_and_fakeid(self):
        """从浏览器获取token和fakeid"""
        # 访问公众号文章管理页面
        self.driver.get(f"{WECHAT_CONFIG['base_url']}/cgi-bin/appmsg")
        
        # 获取当前URL中的token
        current_url = self.driver.current_url
        token = None
        if 'token=' in current_url:
            token = current_url.split('token=')[1].split('&')[0]
        
        # 获取fakeid（需要用户手动选择公众号）
        print("请在浏览器中选择要提取文章的公众号...")
        input("选择完成后按回车键继续...")
        
        # 从URL获取fakeid
        current_url = self.driver.current_url
        fakeid = None
        if 'fakeid=' in current_url:
            fakeid = current_url.split('fakeid=')[1].split('&')[0]
        
        return token, fakeid
    
    def create_excel_with_hyperlinks(self):
        """创建带超链接的Excel文件"""
        print("正在生成Excel文件...")
        
        # 创建DataFrame
        df = pd.DataFrame(self.articles, columns=OUTPUT_CONFIG['columns'])
        
        # 创建Excel文件
        with pd.ExcelWriter(OUTPUT_CONFIG['filename'], engine='openpyxl') as writer:
            df.to_excel(writer, sheet_name=OUTPUT_CONFIG['sheet_name'], index=False)
            
            # 获取工作表
            worksheet = writer.sheets[OUTPUT_CONFIG['sheet_name']]
            
            # 设置列宽
            worksheet.column_dimensions['A'].width = 50  # 标题列
            worksheet.column_dimensions['B'].width = 20  # 时间列
            worksheet.column_dimensions['C'].width = 15  # 地点列
            worksheet.column_dimensions['D'].width = 50  # 链接列
            
            # 为标题添加超链接
            for row_num, article in enumerate(self.articles, start=2):
                cell = worksheet[f'A{row_num}']
                cell.value = f'=HYPERLINK("{article["url"]}", "{article["title"]}")'
                cell.style = 'Hyperlink'
        
        print(f"Excel文件已生成：{OUTPUT_CONFIG['filename']}")
    
    def run(self, username, password):
        """运行提取器"""
        try:
            self.setup_browser()
            
            # 登录
            if not self.login(username, password):
                return False
            
            # 获取token和fakeid
            token, fakeid = self.get_token_and_fakeid()
            
            if not token or not fakeid:
                print("无法获取token或fakeid")
                return False
            
            # 提取文章
            self.extract_all_articles(fakeid, token)
            
            # 生成Excel
            self.create_excel_with_hyperlinks()
            
            return True
            
        except Exception as e:
            print(f"运行出错: {e}")
            return False
        finally:
            if self.driver:
                self.driver.quit()


def main():
    """主函数"""
    print("微信公众号文章信息提取器")
    print("=" * 50)
    
    extractor = WeChatArticleExtractor()
    
    # 获取用户输入
    username = input("请输入微信公众号用户名: ")
    password = input("请输入密码: ")
    
    # 运行提取器
    if extractor.run(username, password):
        print("提取完成！")
    else:
        print("提取失败！")


if __name__ == "__main__":
    main()
